import matplotlib as mpl
from sklearn import preprocessing
from sklearn.metrics import r2_score
from sklearn.preprocessing import PolynomialFeatures

from as_mfs import *
from getData import *

mpl.rcParams['font.sans-serif'] = ['SimHei']
mpl.rcParams['axes.unicode_minus'] = False

current_path = os.path.dirname(__file__)
parent_path = os.path.dirname(current_path)

degree = 4
poly = PolynomialFeatures(degree=degree)


class AN_MFS():
    def __init__(self) -> None:
        self.__init_meta__target()
        self.__init_meta__source()
        self.list_target = load_target_data()
        self.list_source = load_source_data()
        pass

    def __init_meta__source(self) -> None:
        self.vars_name_source = ['alpha', 'CL', 'CD', 'CDp', 'Re', 'Cl/Cd']

        self.x_name_source = ['alpha', 'Re']
        self.x_mean_source = ['攻角', '雷诺数']

        self.y_name_source = ['CD', 'Cl/Cd']
        self.y_mean_source = ['阻力系数', '升阻比']

    def __init_meta__target(self) -> None:
        self.vars_name_target = ['alpha', 'Cl', 'Cd', 'Re', 'Cl/Cd']

        self.x_name_target = ['alpha', 'Re']
        self.x_mean_target = ['攻角', '雷诺数']

        self.y_name_target = ['Cd', 'Cl/Cd']
        self.y_mean_target = ['阻力系数', '升阻比']

    def source_data_init(self):

        self.x_source = self.list_source[0][self.x_name_source].to_numpy()
        self.y_source = self.list_source[0][self.y_name_source].to_numpy()
        for i in range(1, len(self.list_source)):
            self.x_source = np.vstack((self.x_source, self.list_source[i][self.x_name_source].to_numpy()))

        for i in range(1, len(self.list_source)):
            self.y_source = np.vstack((self.y_source, self.list_source[i][self.y_name_source].to_numpy()))

        self.x_source_train = self.x_source[0:self.x_source.shape[0]:2]
        self.x_source_test = self.x_source[1:self.x_source.shape[0]:2]
        self.y_source_train = self.y_source[0:self.y_source.shape[0]:2]
        self.y_source_test = self.y_source[1:self.y_source.shape[0]:2]

        self.x_source_train_st = preprocessing.scale(self.x_source_train)
        self.x_source_test_st = preprocessing.scale(self.x_source_test)

        self.x_source_train_st = poly.fit_transform(self.x_source_train_st)  # 对数据进行预处理，添加一些新的多维特征
        self.x_source_test_st = poly.fit_transform(self.x_source_test_st)

    def target_data_init(self):
        self.x_target = self.list_target[0][self.x_name_target].to_numpy()
        self.y_target = self.list_target[0][self.y_name_target].to_numpy()
        for i in range(1, len(self.list_target)):
            self.x_target = np.vstack((self.x_target, self.list_target[i][self.x_name_target].to_numpy()))

        for i in range(1, len(self.list_target)):
            self.y_target = np.vstack((self.y_target, self.list_target[i][self.y_name_target].to_numpy()))

        # ===========================================================================
        # # 1/2
        # self.title = "1/2风洞数据用于训练"
        # self.x_target_train = self.x_target[0:self.x_target.shape[0]:2]
        # self.x_target_test = self.x_target[1:self.x_target.shape[0]:2]
        # self.y_target_train = self.y_target[0:self.y_target.shape[0]:2]
        # self.y_target_test = self.y_target[1:self.y_target.shape[0]:2]
        # ===========================================================================
        # # 1/3
        # self.title = "1/3风洞数据用于训练"
        # self.x_target_train = self.x_target[0:self.x_target.shape[0]:3]
        # X1 = self.x_target[1:self.x_target.shape[0]:3]
        # X2 = self.x_target[2:self.x_target.shape[0]:3]
        # X3 = np.concatenate((X1, X2))
        # self.x_target_test = X3
        #
        # self.y_target_train = self.y_target[0:self.y_target.shape[0]:3]
        # Y1 = self.y_target[1:self.y_target.shape[0]:3]
        # Y2 = self.y_target[2:self.y_target.shape[0]:3]
        # Y3 = np.concatenate((Y1, Y2))
        # self.y_target_test = Y3
        # ===========================================================================
        # 1/4
        self.title = "1/4风洞数据用于训练"
        self.x_target_train = self.x_target[0:self.x_target.shape[0]:4]
        X1 = self.x_target[1:self.x_target.shape[0]:4]
        X2 = self.x_target[2:self.x_target.shape[0]:4]
        X3 = self.x_target[3:self.x_target.shape[0]:4]
        X4 = np.concatenate((X1, X2, X3))
        self.x_target_test = X4

        self.y_target_train = self.y_target[0:self.y_target.shape[0]:4]
        Y1 = self.y_target[1:self.y_target.shape[0]:4]
        Y2 = self.y_target[2:self.y_target.shape[0]:4]
        Y3 = self.y_target[3:self.y_target.shape[0]:4]
        Y4 = np.concatenate((Y1, Y2, Y3))
        self.y_target_test = Y4
        # ===========================================================================
        # 1/5
        # self.title = "1/5风洞数据用于训练"
        # self.x_target_train = self.x_target[0:self.x_target.shape[0]:5]
        # X1 = self.x_target[1:self.x_target.shape[0]:5]
        # X2 = self.x_target[2:self.x_target.shape[0]:5]
        # X3 = self.x_target[3:self.x_target.shape[0]:5]
        # X4 = self.x_target[4:self.x_target.shape[0]:5]
        # X5 = np.concatenate((X1, X2, X3, X4))
        # self.x_target_test = X5
        #
        # self.y_target_train = self.y_target[0:self.y_target.shape[0]:5]
        #
        # Y1 = self.y_target[1:self.y_target.shape[0]:5]
        # Y2 = self.y_target[2:self.y_target.shape[0]:5]
        # Y3 = self.y_target[3:self.y_target.shape[0]:5]
        # Y4 = self.y_target[4:self.y_target.shape[0]:5]
        # Y5 = np.concatenate((Y1, Y2, Y3, Y4))
        # self.y_target_test = Y5
        # ===========================================================================
        # 1/6
        # self.title = "1/6风洞数据用于训练"
        # self.x_target_train = self.x_target[3:self.x_target.shape[0]:6]
        # X1 = self.x_target[0:self.x_target.shape[0]:6]
        # X2 = self.x_target[1:self.x_target.shape[0]:6]
        # X3 = self.x_target[2:self.x_target.shape[0]:6]
        # X4 = self.x_target[4:self.x_target.shape[0]:6]
        # X5 = self.x_target[5:self.x_target.shape[0]:6]
        # self.x_target_test = np.concatenate((X1, X2, X3, X4, X5))
        #
        # self.y_target_train = self.y_target[3:self.y_target.shape[0]:6]
        #
        # Y1 = self.y_target[0:self.y_target.shape[0]:6]
        # Y2 = self.y_target[1:self.y_target.shape[0]:6]
        # Y3 = self.y_target[2:self.y_target.shape[0]:6]
        # Y4 = self.y_target[4:self.y_target.shape[0]:6]
        # Y5 = self.y_target[5:self.y_target.shape[0]:6]
        # self.y_target_test = np.concatenate((Y1, Y2, Y3, Y4, Y5))
        # ===========================================================================
        # 1/7
        # self.title = "1/7风洞数据用于训练"
        # self.x_target_train = self.x_target[3:self.x_target.shape[0]:7]
        # X1 = self.x_target[0:self.x_target.shape[0]:7]
        # X2 = self.x_target[1:self.x_target.shape[0]:7]
        # X3 = self.x_target[2:self.x_target.shape[0]:7]
        # X4 = self.x_target[4:self.x_target.shape[0]:7]
        # X5 = self.x_target[5:self.x_target.shape[0]:7]
        # X6 = self.x_target[6:self.x_target.shape[0]:7]
        # self.x_target_test = np.concatenate((X1, X2, X3, X4, X5, X6))
        #
        # self.y_target_train = self.y_target[3:self.y_target.shape[0]:7]
        # Y1 = self.y_target[0:self.y_target.shape[0]:7]
        # Y2 = self.y_target[1:self.y_target.shape[0]:7]
        # Y3 = self.y_target[2:self.y_target.shape[0]:7]
        # Y4 = self.y_target[4:self.y_target.shape[0]:7]
        # Y5 = self.y_target[5:self.y_target.shape[0]:7]
        # Y6 = self.y_target[6:self.y_target.shape[0]:7]
        # self.y_target_test = np.concatenate((Y1, Y2, Y3, Y4, Y5, Y6))
        # ===========================================================================
        # 1/8
        # self.title = "1/8风洞数据用于训练"
        # self.x_target_train = self.x_target[0:self.x_target.shape[0]:8]
        # X1 = self.x_target[1:self.x_target.shape[0]:8]
        # X2 = self.x_target[2:self.x_target.shape[0]:8]
        # X3 = self.x_target[3:self.x_target.shape[0]:8]
        # X4 = self.x_target[4:self.x_target.shape[0]:8]
        # X5 = self.x_target[5:self.x_target.shape[0]:8]
        # X6 = self.x_target[6:self.x_target.shape[0]:8]
        # X7 = self.x_target[7:self.x_target.shape[0]:8]
        # self.x_target_test = np.concatenate((X1, X2, X3, X4, X5, X6, X7))
        #
        # self.y_target_train = self.y_target[0:self.y_target.shape[0]:8]
        # Y1 = self.y_target[1:self.y_target.shape[0]:8]
        # Y2 = self.y_target[2:self.y_target.shape[0]:8]
        # Y3 = self.y_target[3:self.y_target.shape[0]:8]
        # Y4 = self.y_target[4:self.y_target.shape[0]:8]
        # Y5 = self.y_target[5:self.y_target.shape[0]:8]
        # Y6 = self.y_target[6:self.y_target.shape[0]:8]
        # Y7 = self.y_target[7:self.y_target.shape[0]:8]
        # self.y_target_test = np.concatenate((Y1, Y2, Y3, Y4, Y5, Y6, Y7))
        # ===========================================================================
        # 1/9
        # self.title = "1/9风洞数据用于训练"
        # self.x_target_train = self.x_target[4:self.x_target.shape[0]:9]
        # X1 = self.x_target[0:self.x_target.shape[0]:9]
        # X2 = self.x_target[1:self.x_target.shape[0]:9]
        # X3 = self.x_target[2:self.x_target.shape[0]:9]
        # X4 = self.x_target[3:self.x_target.shape[0]:9]
        # X5 = self.x_target[5:self.x_target.shape[0]:9]
        # X6 = self.x_target[6:self.x_target.shape[0]:9]
        # X7 = self.x_target[7:self.x_target.shape[0]:9]
        # X8 = self.x_target[8:self.x_target.shape[0]:9]
        # self.x_target_test = np.concatenate((X1, X2, X3, X4, X5, X6, X7, X8))
        #
        # self.y_target_train = self.y_target[4:self.y_target.shape[0]:9]
        # Y1 = self.y_target[0:self.y_target.shape[0]:9]
        # Y2 = self.y_target[1:self.y_target.shape[0]:9]
        # Y3 = self.y_target[2:self.y_target.shape[0]:9]
        # Y4 = self.y_target[3:self.y_target.shape[0]:9]
        # Y5 = self.y_target[5:self.y_target.shape[0]:9]
        # Y6 = self.y_target[6:self.y_target.shape[0]:9]
        # Y7 = self.y_target[7:self.y_target.shape[0]:9]
        # Y8 = self.y_target[8:self.y_target.shape[0]:9]
        # self.y_target_test = np.concatenate((Y1, Y2, Y3, Y4, Y5, Y6, Y7, Y8))
        # ===========================================================================
        # 1/10
        # self.title = "1/10风洞数据用于训练"
        # self.x_target_train = self.x_target[0:self.x_target.shape[0]:10]
        # X1 = self.x_target[1:self.x_target.shape[0]:10]
        # X2 = self.x_target[2:self.x_target.shape[0]:10]
        # X3 = self.x_target[3:self.x_target.shape[0]:10]
        # X4 = self.x_target[4:self.x_target.shape[0]:10]
        # X5 = self.x_target[5:self.x_target.shape[0]:10]
        # X6 = self.x_target[6:self.x_target.shape[0]:10]
        # X7 = self.x_target[7:self.x_target.shape[0]:10]
        # X8 = self.x_target[8:self.x_target.shape[0]:10]
        # X9 = self.x_target[9:self.x_target.shape[0]:10]
        # self.x_target_test = np.concatenate((X1, X2, X3, X4, X5, X6, X7, X8, X9))
        #
        # self.y_target_train = self.y_target[0:self.y_target.shape[0]:10]
        # Y1 = self.y_target[1:self.y_target.shape[0]:10]
        # Y2 = self.y_target[2:self.y_target.shape[0]:10]
        # Y3 = self.y_target[3:self.y_target.shape[0]:10]
        # Y4 = self.y_target[4:self.y_target.shape[0]:10]
        # Y5 = self.y_target[5:self.y_target.shape[0]:10]
        # Y6 = self.y_target[6:self.y_target.shape[0]:10]
        # Y7 = self.y_target[7:self.y_target.shape[0]:10]
        # Y8 = self.y_target[8:self.y_target.shape[0]:10]
        # Y9 = self.y_target[9:self.y_target.shape[0]:10]
        # self.y_target_test = np.concatenate((Y1, Y2, Y3, Y4, Y5, Y6, Y7, Y8, Y9))
        # ===========================================================================

        self.x_target_train_st = preprocessing.scale(self.x_target_train)
        self.x_target_test_st = preprocessing.scale(self.x_target_test)

        self.x_target_train_st = poly.fit_transform(self.x_target_train_st)  # 对数据进行预处理，添加一些新的多维特征
        self.x_target_test_st = poly.fit_transform(self.x_target_test_st)

    def source_model_fit(self):
        self.source_model = as_mfs(self.x_source_train_st, self.y_source_train)

    def train_model_result(self):
        y_pre = self.source_model.predict(self.x_source_test_st)
        # for i in range(len(self.y_name_source)):
        #     print("源域AF_MFS模型", self.y_name_source[i], "MSE误差为", err(y_pre[:, i], self.y_source_test[:, i]))
        #     print("源域AF_MFS模型", self.y_name_source[i], "R2误差为", r2_score(y_pre[:, i], self.y_source_test[:, i]))
        label = "快速工程"
        title = ""
        pre_plt(self.x_source_train, self.y_source_train, self.x_source_test, self.y_source_test, y_pre, label, title)

    def different(self):
        self.d_y = self.y_target_train - self.source_model.predict(self.x_target_train_st)

    def d_model_fit(self):
        self.d_model = as_mfs(self.x_target_train_st, self.d_y)

    def final_result(self):
        y_pre = self.source_model.predict(self.x_target_test_st) + self.d_model.predict(self.x_target_test_st)
        for i in range(len(self.y_name_source)):
            print("源域AF_MFS模型", self.y_name_source[i], "MSE误差为", err(y_pre[:, i], self.y_target_test[:, i]))
            print("源域AF_MFS模型", self.y_name_source[i], "R2误差为", r2_score(y_pre[:, i], self.y_target_test[:, i]))
        label = "风洞"
        pre_plt(self.x_target_train, self.y_target_train, self.x_target_test, self.y_target_test, y_pre, label, self.title)


if __name__ == '__main__':
    model = AN_MFS()
    model.source_data_init()
    model.target_data_init()
    model.source_model_fit()
    model.train_model_result()
    model.different()
    model.d_model_fit()
    model.final_result()
